import os
from pathlib import Path
from moviepy.editor import VideoFileClip, ImageClip, AudioFileClip, concatenate_videoclips, concatenate_audioclips
from mutagen.mp3 import MP3
import re
from PIL import Image
import numpy as np

def natural_sort_key(s):
    """用于实现文件夹的自然排序"""
    return [int(text) if text.isdigit() else text.lower()
            for text in re.split('([0-9]+)', str(s))]

def get_audio_duration(audio_path):
    """获取音频文件的持续时间"""
    try:
        audio = MP3(audio_path)
        return audio.info.length
    except:
        return 5.0  # 默认持续时间为5秒

def process_image(image_path, target_width=1920, target_height=1080):
    """处理图片，包括尺寸调整和优化"""
    # 打开图片
    with Image.open(image_path) as img:
        original_width, original_height = img.size
        
        # 计算缩放比例（总是缩放到目标尺寸）
        width_ratio = target_width / original_width
        height_ratio = target_height / original_height
        
        # 选择较小的缩放比例，确保整个图片都能显示在屏幕内
        scale = min(width_ratio, height_ratio)
        
        # 计算新的尺寸
        new_width = int(original_width * scale)
        new_height = int(original_height * scale)
        
        # 放大图片，使用高质量的重采样方法
        img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
        
        # 创建新的背景图片（黑色背景）
        background = Image.new('RGB', (target_width, target_height), (0, 0, 0))
        
        # 计算粘贴位置（居中）
        paste_x = (target_width - new_width) // 2
        paste_y = (target_height - new_height) // 2
        
        # 将调整后的图片粘贴到背景上
        background.paste(img, (paste_x, paste_y))
        
        # 保存处理后的图片到临时文件
        temp_path = str(image_path).replace('.jpg', '_processed.jpg')
        background.save(temp_path, 'JPEG', quality=95)
        return temp_path

def create_video_batch(folders, batch_number):
    """为一批文件夹创建视频"""
    output_path = f"output_video_batch_{batch_number}.mp4"
    print(f"正在处理第 {batch_number} 批视频，包含 {len(folders)} 个文件夹")
    
    video_clips = []
    temp_files = []  # 用于跟踪临时文件
    
    for folder in folders:
        print(f"处理文件夹: {folder}")
        # 查找图片文件
        image_files = list(folder.glob("panel_*.jpg"))
        if not image_files:
            continue
            
        # 获取第一个图片文件（假设每个文件夹只有一个图片）
        image_path = image_files[0]
        
        # 处理图片
        processed_image_path = process_image(str(image_path))
        if processed_image_path != str(image_path):
            temp_files.append(processed_image_path)
        
        # 查找并排序所有音频文件
        audio_files = sorted(list(folder.glob("*.mp3")), key=natural_sort_key)
        
        if not audio_files:
            # 如果没有音频文件，使用默认时长
            duration = 2.5
            video_clip = ImageClip(processed_image_path, duration=duration)
        else:
            # 创建所有音频片段
            audio_clips = []
            total_duration = 0
            
            for audio_path in audio_files:
                audio_duration = get_audio_duration(audio_path)
                total_duration += audio_duration
                audio_clip = AudioFileClip(str(audio_path))
                audio_clips.append(audio_clip)
            
            # 创建持续时间等于所有音频总时长的视频片段
            video_clip = ImageClip(processed_image_path, duration=total_duration)
            
            # 连接所有音频片段
            if len(audio_clips) > 1:
                final_audio = concatenate_audioclips(audio_clips)
            else:
                final_audio = audio_clips[0]
                
            # 添加音频到视频
            video_clip = video_clip.set_audio(final_audio)
        
        video_clips.append(video_clip)
    
    if video_clips:
        try:
            # 连接所有视频片段
            final_video = concatenate_videoclips(video_clips, method="compose")
            
            # 导出视频，使用更高的比特率和编码器设置
            final_video.write_videofile(
                output_path,
                fps=24,
                codec='libx264',
                audio_codec='aac',
                bitrate='8000k',
                preset='medium',  # 改用medium预设，在质量和速度之间取得平衡
                threads=4,
                audio_bitrate='192k',
                temp_audiofile=f"temp-audio-batch-{batch_number}.m4a",  # 指定临时音频文件
                remove_temp=True  # 处理完成后删除临时文件
            )
            
            # 清理资源
            final_video.close()
            for clip in video_clips:
                clip.close()
        finally:
            # 清理临时文件
            for temp_file in temp_files:
                try:
                    os.remove(temp_file)
                except:
                    pass

def create_videos():
    """创建多个视频，每100个文件夹一个视频，如果最后不足100个则并入前一批"""
    # 设置输入路径
    input_dir = Path("bofan")
    
    # 获取并排序所有子文件夹
    all_folders = sorted([f for f in input_dir.iterdir() if f.is_dir()], key=natural_sort_key)
    total_folders = len(all_folders)
    
    # 计算需要生成的批次数
    batch_size = 100
    
    # 如果总文件夹数小于等于100，直接处理为一个视频
    if total_folders <= batch_size:
        print(f"总共有 {total_folders} 个文件夹，将生成一个视频")
        create_video_batch(all_folders, 1)
        return
    
    # 计算完整的批次数（不包括最后一个不完整的批次）
    full_batches = total_folders // batch_size
    remaining_folders = total_folders % batch_size
    
    # 如果剩余文件夹数量小于50，则并入最后一个完整批次
    if remaining_folders > 0 and remaining_folders < 50:
        num_batches = full_batches
        print(f"总共有 {total_folders} 个文件夹，将分成 {num_batches} 个批次处理")
        print(f"最后 {remaining_folders} 个文件夹将并入最后一批")
    else:
        num_batches = (total_folders + batch_size - 1) // batch_size
        print(f"总共有 {total_folders} 个文件夹，将分成 {num_batches} 个批次处理")
    
    # 按批次处理文件夹
    for batch_num in range(num_batches):
        start_idx = batch_num * batch_size
        
        # 如果是最后一个批次，并且剩余文件夹少于50，则包含所有剩余文件夹
        if batch_num == num_batches - 1 and remaining_folders < 50:
            end_idx = total_folders
        else:
            end_idx = min((batch_num + 1) * batch_size, total_folders)
        
        batch_folders = all_folders[start_idx:end_idx]
        
        print(f"\n开始处理第 {batch_num + 1} 批视频 (文件夹 {start_idx + 1} 到 {end_idx})")
        print(f"本批次包含 {len(batch_folders)} 个文件夹")
        create_video_batch(batch_folders, batch_num + 1)
        print(f"第 {batch_num + 1} 批视频处理完成")

if __name__ == "__main__":
    create_videos() 